#include "compi/compi.h"

#include <iostream>
#include <mpi.h>
#include <sstream>
#include <vector>

using namespace compi;

int
main(int argc, char ** argv)
{
    Environment env(ThreadLevel::Multiple);
    Context & ctx = Context::for_comm(MPI_COMM_WORLD, env);

    const int world_size = ctx.size();
    const int world_rank = ctx.rank();

    // 通过 MPI_Dims_create 推导二维网格
    std::vector<int> dims(2, 0);
    auto dims_res =
        ctx.call_mpi([&]() -> int { return MPI_Dims_create(world_size, static_cast<int>(dims.size()), dims.data()); },
                     "MPI_Dims_create");
    if (dims_res.is_error())
    {
        if (world_rank == 0)
        {
            std::cerr << "MPI_Dims_create failed: " << dims_res.error_message() << std::endl;
        }
        return 1;
    }

    std::vector<int> periods(dims.size(), 0);
    topology::CartCreateParams params;
    params.dims = dims;
    params.periods = periods;
    params.reorder = true;

    auto topo_res = topology::create_cartesian(ctx, params);
    if (topo_res.is_error())
    {
        if (world_rank == 0)
        {
            std::cerr << "Failed to create cartesian topology: " << topo_res.error_message() << std::endl;
        }
        return 1;
    }

    topology::TopologyHandle topo = std::move(topo_res).value();
    if (!topo.valid())
    {
        if (world_rank == 0)
        {
            std::cout << "Process excluded from cartesian communicator." << std::endl;
        }
        return 0;
    }

    auto coords_res = topology::cart_coords(topo, world_rank);
    if (coords_res.is_error())
    {
        if (world_rank == 0)
        {
            std::cerr << "cart_coords failed: " << coords_res.error_message() << std::endl;
        }
        return 1;
    }
    const auto coords = coords_res.value();

    auto rerank_res = topology::cart_rank(topo, coords);
    if (rerank_res.is_error())
    {
        if (world_rank == 0)
        {
            std::cerr << "cart_rank failed: " << rerank_res.error_message() << std::endl;
        }
        return 1;
    }
    const int cart_rank = rerank_res.value();

    // 打印拓扑信息
    std::ostringstream oss;
    oss << "World rank " << world_rank << "/" << world_size << " -> cart rank " << cart_rank << " coords [";
    for (std::size_t i = 0; i < coords.size(); ++i)
    {
        oss << coords[i];
        if (i + 1 < coords.size())
            oss << ", ";
    }
    oss << "]";

    std::cout << oss.str() << std::endl;

    return 0;
}
